import numpy as np
import matplotlib.pyplot as plt
# 在程序开头添加
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示问题
import matplotlib


# 设置中文字体支持
def set_chinese_font():
    """设置中文字体支持"""
    try:
        # 尝试使用系统中常见的中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans Fallback']
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    except:
        print("警告: 中文字体设置失败，图表中的中文可能无法正常显示")


def linear_regression(M, U):
    """
    使用最小二乘法进行线性回归分析

    参数:
    M: 质量数据列表
    U: 电压数据列表

    返回:
    B: 斜率
    C: 截距
    r: 相关系数
    """

    # 将输入转换为numpy数组
    M = np.array(M, dtype=float)
    U = np.array(U, dtype=float)

    # 检查数据长度是否一致
    if len(M) != len(U):
        raise ValueError("Mass data M and voltage data U must have the same length")

    n = len(M)

    # 计算各项和
    sum_U = np.sum(U)
    sum_M = np.sum(M)
    sum_U2 = np.sum(U ** 2)
    sum_M2 = np.sum(M ** 2)
    sum_UM = np.sum(U * M)

    # 计算斜率B和截距C
    denominator = n * sum_U2 - sum_U ** 2
    if denominator == 0:
        raise ValueError("Data points are collinear, cannot calculate slope")

    B = (n * sum_UM - sum_U * sum_M) / denominator
    C = (sum_M - B * sum_U) / n

    # 计算相关系数r
    numerator_r = n * sum_UM - sum_U * sum_M
    denominator_r = np.sqrt((n * sum_U2 - sum_U ** 2) * (n * sum_M2 - sum_M ** 2))

    if denominator_r == 0:
        r = 0
    else:
        r = numerator_r / denominator_r

    return B, C, r


def plot_regression(M, U, B, C, r):
    """
    绘制回归直线和原始数据点
    """
    plt.figure(figsize=(10, 6))

    # 绘制原始数据点
    plt.scatter(U, M, color='blue', label='Data points', s=50)

    # 绘制回归直线
    U_fit = np.linspace(min(U), max(U), 100)
    M_fit = B * U_fit + C
    plt.plot(U_fit, M_fit, color='red', linewidth=2,
             label=f'Regression line: M = {B:.4f}U + {C:.4f}')

    plt.xlabel('Voltage U', fontsize=12)
    plt.ylabel('Mass M', fontsize=12)
    plt.title(f'Linear Regression Analysis (Correlation coefficient r = {r:.6f})', fontsize=14)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def main():
    # 设置中文字体
    set_chinese_font()

    print("=== Least Squares Linear Regression Analysis ===")
    print("Please input mass M and voltage U data")
    print("Format: Input multiple data points separated by spaces")

    try:
        # 获取输入数据
        M_input = input("Enter mass M data (separated by spaces): ").split()
        U_input = input("Enter voltage U data (separated by spaces): ").split()

        # 转换为浮点数
        M = [float(x) for x in M_input]
        U = [float(x) for x in U_input]

        # 进行线性回归分析
        B, C, r = linear_regression(M, U)

        # 输出结果
        print("\n" + "=" * 50)
        print("Linear Regression Results:")
        print(f"Slope B = {B:.6f}")
        print(f"Intercept C = {C:.6f}")
        print(f"Correlation coefficient r = {r:.6f}")
        print(f"Regression equation: M = {B:.6f}U + {C:.6f}")

        # 计算拟合值
        print("\nData fitting results:")
        print("Voltage U\tActual Mass M\tFitted Mass M_fit\tResidual")
        for i in range(len(M)):
            M_fit = B * U[i] + C
            residual = M[i] - M_fit
            print(f"{U[i]:.2f}\t\t{M[i]:.6f}\t{M_fit:.6f}\t\t{residual:.6f}")

        # 绘制图形
        plot_regression(M, U, B, C, r)

    except ValueError as e:
        print(f"Input error: {e}")
    except Exception as e:
        print(f"Error during calculation: {e}")


# 使用numpy的现成函数进行验证
def numpy_verification(M, U):
    """
    使用numpy的polyfit函数进行验证
    """
    M = np.array(M)
    U = np.array(U)

    # 使用numpy进行线性拟合
    coefficients = np.polyfit(U, M, 1)
    B_np = coefficients[0]
    C_np = coefficients[1]

    # 计算相关系数
    r_np = np.corrcoef(U, M)[0, 1]

    return B_np, C_np, r_np


# 替代方案：如果字体问题仍然存在，使用英文版本
def english_version():
    """完全使用英文的版本"""
    print("=== Least Squares Linear Regression Analysis ===")
    print("Please input mass M and voltage U data")
    print("Format: Input multiple data points separated by spaces")

    try:
        M_input = input("Enter mass M data (separated by spaces): ").split()
        U_input = input("Enter voltage U data (separated by spaces): ").split()

        M = [float(x) for x in M_input]
        U = [float(x) for x in U_input]

        B, C, r = linear_regression(M, U)

        print("\n" + "=" * 50)
        print("Linear Regression Results:")
        print(f"Slope B = {B:.6f}")
        print(f"Intercept C = {C:.6f}")
        print(f"Correlation coefficient r = {r:.6f}")
        print(f"Regression equation: M = {B:.6f}U + {C:.6f}")

        print("\nData fitting results:")
        print("U\tM_actual\tM_fit\tResidual")
        for i in range(len(M)):
            M_fit = B * U[i] + C
            residual = M[i] - M_fit
            print(f"{U[i]:.2f}\t{M[i]:.6f}\t{M_fit:.6f}\t{residual:.6f}")

        # 绘制图形（英文标签）
        plt.figure(figsize=(10, 6))
        plt.scatter(U, M, color='blue', label='Data points', s=50)
        U_fit = np.linspace(min(U), max(U), 100)
        M_fit = B * U_fit + C
        plt.plot(U_fit, M_fit, color='red', linewidth=2,
                 label=f'Regression: M = {B:.4f}U + {C:.4f}')
        plt.xlabel('Voltage U')
        plt.ylabel('Mass M')
        plt.title(f'Linear Regression (r = {r:.6f})')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    except Exception as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    # 尝试使用中文版本，如果失败则使用英文版本
    try:
        main()
    except Exception as e:
        print(f"Chinese version failed: {e}")
        print("Switching to English version...")
        english_version()

    # 示例数据测试
    print("\n" + "=" * 50)
    print("Example data test:")
    example_M = [1.2, 2.1, 3.0, 3.9, 4.8]
    example_U = [1.0, 2.0, 3.0, 4.0, 5.0]
    print(f"Example mass M: {example_M}")
    print(f"Example voltage U: {example_U}")

    B_ex, C_ex, r_ex = linear_regression(example_M, example_U)
    print(f"Slope B = {B_ex:.6f}")
    print(f"Intercept C = {C_ex:.6f}")
    print(f"Correlation coefficient r = {r_ex:.6f}")